"""
用于从gazebo生成的仿真图像中获取GT信息，边缘点，无拟合，无内点
"""

import numpy as np
import cv2
import os
import tqdm
import argparse
import yaml
import yaml.parser
from gt_gen.base import ContourGT
from typing import Iterable

LATITUDE_RANGE = 45.2410 - 42.8613  # deg
LONGITUDE_RANGE = -51.6879 - (-52.0974)  # deg
X_PIXEL_RANGE = 48131
Y_PIXEL_RANGE = 8264
RATE = 2
raw_reso = 1.5
scale_factor = raw_reso * RATE / 0.2
BASE_BATCH = 1000


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )
    arg_parser.add_argument(
        "-l",
        "--label_dir",
        type=str,
        default="/disk527/sdb1/a804_cbf/datasets/lunar_crater",
        help="The directory of the label data",
    )
    arg_parser.add_argument("--cache", action="store_true", help="Whether to use cache")
    arg_parser.add_argument(
        "--image", action="store_true", help="Whether to use image as background"
    )
    return arg_parser.parse_args()


class ChangE1_5m(ContourGT):
    def __init__(
        self,
        label_dir,
        origin_point: tuple,
        K: tuple,
        size: tuple,
        rate: int,
        *,
        cache_dir=None,
        modifier,
        lat: tuple,
        lon: tuple,
        pixel_range: tuple,
        row: list,
        col: list,
    ):
        super().__init__(
            label_dir,
            origin_point,
            K,
            size,
            cache_dir=cache_dir,
            modifier=modifier,
            lat=lat,
            lon=lon,
            pixel_range=pixel_range,
            preload={"row": row, "col": col, "rate": rate},
            findCraters={"row0": row[0], "col0": col[0]},
        )

    def findCraters(self, dem, data, idx, row0, col0):
        """
        从DEM中找到所有的陨石坑的边缘三维点，并将其坐标转为直角坐标
        """
        points = []
        points_id = []
        # temp = (dem - dem.min()) / (dem.max() - dem.min()) * 255
        for id, (x, y, lat, lon, d) in tqdm.tqdm(
            enumerate(data), desc="Preparing for crater points", total=len(idx)
        ):
            # 把所有完整弧的陨石坑都包含进来
            # if not self.catalog(idx[id]):
            #     continue
            D_x, D_y = self.transform.distortCircle(lat, d * 1000)
            mask = np.zeros_like(dem, dtype=np.int16)
            # # 画边缘的椭圆
            cv2.ellipse(
                mask,
                (round(y) - col0 * BASE_BATCH, round(x) - row0 * BASE_BATCH),
                (round(D_y / 2), round(D_x / 2)),
                0,
                0,
                360,
                1,
                2,
            )
            # cv2.ellipse(
            #     temp,
            #     (round(y) - col0 * BASE_BATCH, round(x) - row0 * BASE_BATCH),
            #     (round(D_y / 2), round(D_x / 2)),
            #     0,
            #     0,
            #     360,
            #     0,
            #     2,
            # )
            # 画中心点
            # cv2.circle(mask, (int(y), int(x)), 1, 1)
            img = mask * dem
            if (mask != 0).sum() == 0:
                continue
            dr = img[mask != 0]
            coor = cv2.findNonZero(mask).squeeze()
            lat, lon = self.transform.plane2sphere(
                coor[:, 1] + row0 * BASE_BATCH, coor[:, 0] + col0 * BASE_BATCH
            )
            ids = np.ones_like(lat, dtype=np.int32) * id
            points.append(self.transform.sphere2xyz(lat, lon, dr))
            points_id.append(ids)
        # 原点平移至标准位置
        # cv2.imwrite(f"test.png", temp)
        points_ = np.concatenate(points, axis=1)
        points_id = np.concatenate(points_id)
        R = self.transform.rotate(*self.origin_point)
        T = self.transform.sphere2xyz(*self.origin_point, 0)
        T = -R @ np.array(T)
        points_ = np.pad(points_, ((0, 1), (0, 0)), "constant", constant_values=1)
        H = np.hstack((R, T[:, None]))
        H = np.vstack((H, np.array([0, 0, 0, 1])))
        cali_points = H @ points_
        return cali_points, points_id

    def preload(self, label_dir, *, row: Iterable, col: Iterable, rate: int):
        def _load(label_dir, row_, col_):
            dem = np.load(
                os.path.join(
                    label_dir,
                    f"dem_int16/{row_}/chang_e_{row_}_{col_}_{BASE_BATCH}.npz",
                )
            )["image"]
            data = []
            idx = []

            if os.path.exists(
                os.path.join(
                    label_dir,
                    f"labels/circle/{row_}/chang_e_{row_}_{col_}_{BASE_BATCH}.txt",
                )
            ):
                with open(
                    os.path.join(
                        label_dir,
                        f"labels/circle/{row_}/chang_e_{row_}_{col_}_{BASE_BATCH}.txt",
                    ),
                    "r",
                ) as reader:
                    for x in map(lambda line: line.strip().split(","), reader):
                        data.append(
                            [
                                float(x[0]) + row_ * BASE_BATCH,  # x
                                float(x[1]) + col_ * BASE_BATCH,  # y
                                float(x[5]),  # lat
                                float(x[6]),  # lon
                                float(x[7]),  # radius
                            ]
                        )
                        idx.append(x[4])

                return dem, data, idx

        dem = np.zeros(
            (BASE_BATCH * len(row), BASE_BATCH * len(col)),
            dtype=np.int16,
        )
        data = []
        idx = []
        for row_ in row:
            for col_ in col:
                dem_, data_, idx_ = _load(label_dir, row_, col_)
                dem[
                    (row_ - row[0]) * BASE_BATCH : (row_ + 1 - row[0]) * BASE_BATCH,
                    (col_ - col[0]) * BASE_BATCH : (col_ + 1 - col[0]) * BASE_BATCH,
                ] = dem_[:-rate, :-rate]
                if len(data_) != 0:
                    data.append(data_)
                    idx.append(idx_)
        data = np.concatenate(data)
        idx = np.concatenate(idx)
        return dem, data, idx


def gt_contour(sub_dir, collected_dir, cache_dir):
    with open(os.path.join(collected_dir, "config.yaml"), "r") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    row = config["world"]["row"]
    col = config["world"]["col"]
    num_batches = X_PIXEL_RANGE // BASE_BATCH
    num_samples = Y_PIXEL_RANGE // BASE_BATCH
    delta_theta = LATITUDE_RANGE / num_batches
    delta_phi = LONGITUDE_RANGE / num_samples
    origin_point = (
        45.2410 - delta_theta * row[len(row) // 2],
        -52.0974 + delta_phi * col[len(col) // 2],
    )

    gt_gen = ChangE1_5m(
        sub_dir,
        origin_point,
        config["camera"]["P"],
        (config["camera"]["height"], config["camera"]["width"]),
        RATE,
        cache_dir=cache_dir,
        modifier=scale_factor,
        lat=(45.2410, 42.8613),
        lon=(-52.0974, -51.6879),
        pixel_range=(X_PIXEL_RANGE, Y_PIXEL_RANGE),
        row=row,
        col=col,
    )
    image_names = list(
        map(
            lambda x: float(x.removesuffix(".png")),
            os.listdir(f"{collected_dir}/images"),
        )
    )
    image_names.sort()
    it = 0
    image_name = image_names.pop(0)
    with open(f"{collected_dir}/pose.csv", "r") as f:
        f.readline()  # 去除标题
        with open(f"{collected_dir}/gt_labels.txt", "w") as writer:
            end_of_file = False
            for line in tqdm.tqdm(
                f, desc="Generating GT images", total=len(image_names)
            ):
                # 某个时刻的姿态
                data = line.split(",")
                time = float(data[0])
                if image_name - time > 2e-4:
                    continue
                else:
                    time = f"{image_name:.6f}"
                    if len(image_names) != 0:
                        image_name = image_names.pop(0)
                    else:
                        end_of_file = True

                pose = list(float(x) for x in data[1:4])
                ori = list(float(x) for x in data[4:])
                img = cv2.imread(f"{collected_dir}/images/{time}.png")

                if np.std(img) < 2:
                    continue
                flag = gt_gen(
                    time, pose, ori, collected_dir, img if args.image else None
                )
                if flag:
                    print(f"{time}", file=writer, flush=True)
                    it += 1
                if end_of_file:
                    break
            print("=" * 20 + f"\ntotal labels: {it}\n" + "=" * 20)


if __name__ == "__main__":
    args = arg_parse()
    for collect_dir in args.collect_dir.split(","):
        gt_contour(args.label_dir, collect_dir, "." if args.cache else None)
